#!/usr/bin/python
#coding:utf-8
'''
卷积神经网络-ResNet实例，残差连接结构，多分类器
'''
import tensorflow as tf 
import os 
import _pickle as cPickle 
import numpy as np 

CIFAR_DIR="./cifar-10-batches-py"

def load_data(filename):
    '''
    从文件读取数据
    '''
    with open(filename,'rb') as f:
        data=cPickle.load(f,encoding='bytes')
        return data[b'data'],data[b'labels']

class CifarData:
    '''
    数据处理类
    '''
    def __init__(self,fileList,need_shuffle): #need_shuffle 是否需要shuffle
        all_data=[]
        all_labelList = []
        for filename in fileList:
            data,lableList= load_data(filename)
            
            #获取全量十个类数据
            all_data.append(data)
            all_labelList.append(lableList)

            #获取两类数据
            # for item,label in zip(data,lableList):
            #     if label in [0,1]:
            #         all_data.append(item)
            #         all_labelList.append(label)
        
        self._data= np.vstack(all_data)  #纵行合并为矩阵
        self._data = self._data /127.5 -1 #数据缩放 增加归一化
        self._labels= np.hstack(all_labelList) #横向合并为一维向量
        self._num_example=self._data.shape[0]
        self._need_shuffle = need_shuffle
        self._indicator = 0

        if self._need_shuffle:
            self._shuffle_data()
    
    def _shuffle_data(self):
        '''
        打乱数据
        '''
        p = np.random.permutation(self._num_example)  #实现数据混排
        self._data=self._data[p]
        self._labels=self._labels[p]
    
    def nex_batche(self,batch_size):
        '''
        每次返回batch_size个样本
        数据获取一遍后，shuffle一遍，重新获取
        '''
        end_indicator= self._indicator + batch_size
        if end_indicator > self._num_example:
            if self._need_shuffle:
                self._shuffle_data()
                self._indicator=0
                end_indicator = batch_size
            else:
                raise Exception('have no more examples')
        
        if end_indicator > self._num_example:
            raise Exception("batch size is larger than all example")

        batch_data=self._data[self._indicator:end_indicator]
        batch_labels = self._labels[self._indicator: end_indicator]
        self._indicator=end_indicator

        return batch_data,batch_labels

train_filenames=[os.path.join(CIFAR_DIR,'data_batch_%d'% i) for i in range(1,6)]
test_filenames =[os.path.join(CIFAR_DIR,'test_batch')]

train_data=CifarData(train_filenames,True)
test_data=CifarData(test_filenames,False)


def residual_block(x,output_channel):
    '''
    残差连接基本结构
    args:
        x：输入
        output_channel：输出通道数
    '''
    input_channel= x.get_shape().as_list()[-1]
    if input_channel * 2==output_channel:
        increase_dim=True  #执行降采样
        strides= (2,2) #重置步长
    elif input_channel == output_channel:
        increase_dim= False
        strides=(1,1)
    
    else:
        raise Exception("input channel can`t match output channel")
    conv1 = tf.layers.conv2d(x,
                            output_channel,
                            (3,3),
                            strides=strides,
                            padding="same",
                            activation=tf.nn.relu,
                            name="conv1")
    
    conv2 = tf.layers.conv2d(conv1,
                            output_channel,
                            (3,3),
                            strides=(1,1),
                            padding="same",
                            activation=tf.nn.relu,
                            name='conv2')
    
    #池化
    if increase_dim:
        #[None,image_width,image_height,channel] -> [,,,channel*2]
        pooled_x = tf.layers.average_pooling2d(x,
                                                (2,2),
                                                (2,2),
                                                padding='valid')
        padded_x = tf.pad(pooled_x,[[0,0],
                                    [0,0],
                                    [0,0],
                                    [input_channel//2,input_channel//2]])
    else:
        padded_x=x;
    output_x =  conv2 + padded_x
    return output_x

def res_net(x,
            num_residual_blocks,
            num_filter_base,
            class_num):
    '''
    残差连接实现
    args:
        x:输入
        num_residual_blocks:残差连接块，eg:类型是列表
        num_filter_base:最初通道数目
        class_num:类别数
    '''
    num_subsampling = len(num_residual_blocks)
    layers=[]
    #x: [None,width,height,channel] -> [width, heigth,channel]
    input_size = x.get_shape().as_list()[1:]
    with tf.variable_scope('conv0'):
        conv0 = tf.layers.conv2d(x,
                                num_filter_base,
                                (3,3),
                                strides=(1,1),
                                padding='same',
                                activation=tf.nn.relu,
                                name='conv0')
        layers.append(conv0)
    
    #假设 num_subsampling=4,sample_id =[0,1,2,3]
    for sample_id in range(num_subsampling):
        for i in range(num_residual_blocks[sample_id]):
            with tf.variable_scope('conv%d_%d' % (sample_id,i)):
                conv = residual_block(
                    layers[-1],
                    num_filter_base* (2 ** sample_id))
                layers.append(conv)
    
    multiplier = 2** (num_subsampling - 1)
    assert layers[-1].get_shape().as_list()[1:] == [input_size[0]/multiplier,input_size[1]/multiplier,num_filter_base * multiplier]
    
    with tf.variable_scope("fc"):
        #layer[-1].shape: [None,width,heigth,channel]
        global_pool = tf.reduce_mean(layers[-1],[1,2])
        logits = tf.layers.dense(global_pool,class_num)
        layers.append(logits)

    return layers[-1]


#构建计算图
x= tf.placeholder(tf.float32,[None,3072]) #placeholder(占位符
y= tf.placeholder(tf.int64,[None]) #None 样本数不确定

#数据转换
x_image=tf.reshape(x,[-1,3,32,32])   # 格式转换
x_image=tf.transpose(x_image,perm=[0,2,3,1])  #通道转换，编程正确顺序

y_= res_net(x_image,[2,3,2],32,10)

###########平方差损失函数###########
#针对当前数据集，该损失函验证集成功率40%
#api: e^x / sum(e^x)
# p_y= tf.nn.softmax(y_)
# #5 --> [0000100000]
# y_one_hot= tf.one_hot(y,10,dtype=tf.float32)
# loss=tf.reduce_mean(tf.square(y_one_hot - p_y))  #reduce_mean 求均值  square：求平方，结果为真实值和预测值差的平方

###########交叉熵损失函数##########
#针对当前数据集，损失函数验证集成功率30%
loss=tf.losses.sparse_softmax_cross_entropy(labels=y,logits=y_)

'''
#sigmoid 
#[None,10]
p_y_1 = tf.nn.sigmoid(y_)  #传入sigmoid得到概率值
#[None,1]
y_reshaped=tf.reshape(y,(-1,1))  #数据类型 int64
y_reshaped_float = tf.cast(y_reshaped,tf.float32)  #类型转换 int64 转 float32

#平方差计算损失函数
loss=tf.reduce_mean(tf.square(y_reshaped_float - p_y_1))  #reduce_mean 求均值  square：求平方，结果为真实值和预测值差的平方
'''

#准确率
predict =tf.argmax(y_,1) 

#[1,0,1,1...]
correct_prediction= tf.equal(predict,y) #得到正确预测值
accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float64)) #得到准确率

#梯度下降方法
with tf.name_scope('train_op'):
    train_op = tf.train.AdamOptimizer(1e-3).minimize(loss)  #1e-3：10的负三次方

#创建绘画
init = tf.global_variables_initializer()

batch_size=20
train_steps=100000
test_steps=100

with tf.Session() as sess:
    sess.run(init)
    for i in range(train_steps):
        batch_data,batch_labels=train_data.nex_batche(batch_size)
        loss_val,acc_val,_=sess.run(
                    [loss,accuracy,train_op],
                    feed_dict={x:batch_data, y:batch_labels})  #加train_op 表示训练，不加表示计算，feed_dict:输入数据
        
        if i % 500 ==0:
            print("[Train] Step: %d, loss: %4.5f, acc: %4.5f" %(i,loss_val,acc_val))
        
        #每5000次test数据集评测
        if (i+1) % 5000 ==0:
            test_data=CifarData(test_filenames,False)
            all_test_acc_val=[]
            for j in range(test_steps):
                test_batch_data,test_batch_labels=test_data.nex_batche(batch_size)
                test_acc_val=sess.run(
                    [accuracy],
                    feed_dict={
                        x:test_batch_data,
                        y:test_batch_labels
                    })
                all_test_acc_val.append(test_acc_val)
            
            test_acc=np.mean(all_test_acc_val)
            print('[Test ] Step: %d, acc: %4.5f' % (i+1,test_acc))